'''
Created on Jun 3, 2010

@author: oabalbin
'''

import os
import sys
import glob
import shutil
import copy
import numpy as np
import scipy as sp
from scipy import stats
import pickle
import operator
from optparse import OptionParser
from datetime import datetime
from collections import deque, defaultdict



# My libraries
import signatures.parsers.read_rnaseq_data as pm
import signatures.parsers.read_microarray_data as pmd
import signatures.parsers.read_gene_lists as gpm
import signatures.db.tables as dt
import signatures.db.query as dq
import signatures.db.msigdb as msigdb
import signatures.common.classes as cm #record, arrayMatrix
import signatures.preprocess.tools as tl #filter_constant_genes, get_variable_genes, 
import signatures.seedgenes.seedlist as sd #get_seed_genes, write_varin_file
import signatures.bfrm.paramfile as bp
import signatures.bfrm.bfrm_run as br
import signatures.bfrm.output_parser as bo
import signatures.plotting.plots as myplt
import signatures.db.gene_annotations as ga
from signatures.stats.Rstats import calc_matrix_correlation, calc_hypergeometric_test, plot_correlation_matrix2
from signatures.plotting.plots import plot_correlation_matrix_py, plot_heatmap_mod2
from signatures.stats.stats import calculate_FDR




class tu_enriched_concepts:
    def __init__(self,thisname,pqvalues_array, concept_names):
        self.name = thisname
        self.pqvalues = pqvalues_array
        self.concept_names = concept_names



def parse_input_annotfiles(annotInputFile):

    expr_file, seed_file, outdirs, db_file, samples, bfrm, param, typefile={},{},{},{},{},{},{},'nofile'

    for line in annotInputFile:
        fields = line.strip('\n')
             
        if fields[0][0] == '' or fields[0][0] == '#':
            continue
        
        if fields[0][0] == '@':
            typefile = fields
            continue
        
        fields = fields.split('=')
        dict = {'@Array expression':expr_file, '@Seed genes':seed_file,'@Database file':db_file,'@Samples':samples,'@Output Directories':outdirs,\
                '@BFRM':bfrm,'@Parameters':param}[typefile]
        if typefile=='@Parameters':
            dict[fields[0]] = fields[1].split(',')
        else:
            dict[fields[0]] = fields[1]
        
        
    return expr_file, seed_file, db_file, samples, bfrm, outdirs, param


def set_gene_list(genelist):
    """
    It creates a unique list of genes.
    """
    newgenelist = deque()
    for g in genelist:
        if g not in newgenelist:
            newgenelist.append(g)
            
    return newgenelist


def get_gene_transcript_association_threshold(correlation_matrix, threshold, tunames, genenames, ezh2_target_genes, outfile1,outfile2, comparison):
    """
    It returns a toptus number of correlated tus with each of the genes in the matrix
    It sorts the the gene-TUs correlation matrix over the columns(TUs) and select the ones
    with the toptus
    
    """    
    tus_indicators = np.array(range(correlation_matrix.shape[1]))
    gene_indicators = np.array(range(correlation_matrix.shape[0]))
    highly_correlated_genes=[]
    TUsDict=defaultdict(list)
    
    if len(gene_indicators) != len(genenames):
        print "The size genename array and the number of genes in the correlation matrix is different"
        print len(gene_indicators), len(genenames)
        sys.exit(0)
    if len(tus_indicators) != len(tunames):
        print "The size tuname array and the number of tus in the correlation matrix is different"
        print len(tus_indicators), len(tunames)
        sys.exit(0)

    
    if comparison =="less":
        toptus_correlation_matrix = correlation_matrix < threshold
    elif comparison =="greater":
        toptus_correlation_matrix = correlation_matrix > threshold
        
    frequency_matrix = np.zeros((correlation_matrix.shape[0],correlation_matrix.shape[1]))
    
    
    for i in range(correlation_matrix.shape[0]):
        # TUs with correlation < than threshold
        thistus = tus_indicators[toptus_correlation_matrix[i,:]==True]
        if len(thistus) > 0: 
            for tu in thistus:
                TUsDict[tu].append(i)
                highly_correlated_genes.append(i)
                
        frequency_matrix[i,thistus] += 1
            
    return  TUsDict, highly_correlated_genes

def get_gene_transcript_association(correlation_matrix, toptus, outfile, comparison):
    """
    It returns a toptus number of correlated tus with each of the genes in the matrix
    It sorts the the gene-TUs correlation matrix over the columns(TUs) and select the ones
    with the toptus
    
    """
    
    toptus=toptus+1
    indicators = np.array(range(correlation_matrix.shape[1]))
    tu_inddict = defaultdict(list)
    all_tus = []
    print 'sorting the correlation matrix ' 
    index_corr_sort_matrix = np.argsort(correlation_matrix, axis=1)
    geneTU_indicator_matrix = np.zeros((correlation_matrix.shape[0],toptus))
    
    toptus_correlation_matrix = np.zeros((correlation_matrix.shape[0],toptus))
    
    for i in range(correlation_matrix.shape[0]):
        if comparison== "less":
            thistus = indicators[index_corr_sort_matrix[i,0:toptus]]
        elif comparison== "greater":
            thistus = indicators[index_corr_sort_matrix[i,-toptus:]]

        geneTU_indicator_matrix[i,:] = thistus

        for tu in thistus:
            tu_inddict[i].append(tu)
            all_tus.append(tu)
        #[tu_inddict[i].append(tu) for tu in thistus] 
        
        
        #all_tus.append(tu)
        # get the correlation matrix for the TUs that become selected as the toptus
        toptus_correlation_matrix[i,:] = correlation_matrix[i,thistus]
        
    #write_array2file(outfile, toptus_correlation_matrix)
        
    return geneTU_indicator_matrix, tu_inddict, all_tus

def get_common_targets(geneset1, geneset2):
    """
    """
    common_genes = geneset1.intersection(geneset2)
    return common_genes

def get_TUs_double_correlated(TUsDictL,TUsDictG, TUs_of_interest):
    """
    """
    TUs_of_interst_dict=defaultdict(list)
    TUs_of_interst_dict_len=defaultdict()
    for tui in TUs_of_interest:
        TUs_of_interst_dict_len[tui] = [len(TUsDictL[tui]), len(TUsDictG[tui])]
        TUs_of_interst_dict[tui].append(TUsDictL[tui])
        TUs_of_interst_dict[tui].append(TUsDictG[tui])
        
    return TUs_of_interst_dict, TUs_of_interst_dict_len


def get_TUs_double_correlated_hugo(TUs_of_interst_dict,noveltus,target_genes,outfile):
    """
    It return the TUs of interest dictionary but using the hugo notation.
    It eliminates the gene isoforms because usually the have the same hugo name
    """
    
    TUs_of_interest_hugo = defaultdict()
    TUs_of_interest_len_hugo=defaultdict()
    for tui, myvals in TUs_of_interst_dict.iteritems():
        mykey = noveltus[tui]
        el1 = list(set(tga.translate_refflat2hugo( target_genes[myvals[0]] )))
        el2 = list(set(tga.translate_refflat2hugo( target_genes[myvals[1]] )))
        #el1, el2 = map(str,myvals[0]), map(str,myvals[1])
        #rpg = TUs_of_interst_dict_len[tui][0]
        #upg = TUs_of_interst_dict_len[tui][1]
        rpg = len(el1)
        upg = len(el2)
        TUs_of_interest_hugo[tui] = [rpg,upg,el1,el2]
        TUs_of_interest_len_hugo=[rpg,upg]
        
        outfile.write(mykey+'\t'+str(rpg)+'\t'+str(upg)+'\t'+",".join(el1).replace(',','\t')+"\t&&&\t"+",".join(el2).replace(',','\t')+'\n')
    
    return TUs_of_interest_hugo, TUs_of_interest_len_hugo


if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--annotFile", dest="annotFile",
                            help="annotation file for all files to use")
    
    (options, args) = optionparser.parse_args()
    
    expr_folder, seed_file, db_file, samples, bfrm, outdirs, run_param = parse_input_annotfiles(open(options.annotFile))
    
    odir = cm.directoryTree()
    tstampedfolder, tstamp = odir.create_tstampdir(outdirs['outdir'])
    dirs = odir.make_directory_tree()


    outfile1 = '/exds/users/oabalbin/projects/rnaseq_data/noveltu_expression_data/test/allgenes_v18/expMat_bfrm_coding_allgenes.txt' #dirs[3]+'/'+tstamp+'_expMat_bfrm_coding.txt'   
    outfile2 = '/exds/users/oabalbin/projects/rnaseq_data/noveltu_expression_data/test/allgenes_v18/expMat_bfrm_noveltu_allgenes.txt' #dirs[3]+'/'+tstamp+'_expMat_bfrm_noveltu.txt'   
    outfile3 = '/exds/users/oabalbin/projects/rnaseq_data/noveltu_expression_data/test/allgenes_v18/correlation_test_allgenes.txt' #dirs[3]+'/'+tstamp+'_correlation_test.txt'
    outfile4 = '/exds/users/oabalbin/projects/rnaseq_data/noveltu_expression_data/test/allgenes_v18/correlation_test2_allgenes.txt' #dirs[3]+'/'+tstamp+'_correlation_test2.txt'
    outfile5 = dirs[3]+'/'+tstamp+'_toptus_correlation_matrix.txt'
    outfile5b = dirs[3]+'/'+tstamp+'_toptus_frequency.txt'
    outfile5c = dirs[3]+'/'+tstamp+'_topgenes_frequency.txt'
    outfile8 = dirs[3]+'/'+tstamp+'_toptus_geneslist.txt'
    outfile6 = '/exds/users/oabalbin/projects/rnaseq_data/noveltu_expression_data/test/frequency_dump_test.txt' #dirs[3]+'/'+tstamp+'_frequency_toptus_correlation_matrix.txt'
    outfile7 = '/exds/users/oabalbin/projects/rnaseq_data/noveltu_expression_data/test/correlation_test3_allgenes.txt'
    concepts = '/exds/users/oabalbin/downloads/msigdb.v2.5.symbols.dictionary'
    
    outfile8 = open(outfile8,'w')
    
    tm = pm.parser_rnaseq()
    tp = pmd.parser()
    #pipe.mt.create_connector()
    gtp = gpm.geneparser()
    tga = ga.gene_annotation("localhost", "oabalbin", "oscar", "gene_annotation")
    protgenes=True

    ################ The script
        
    print 'Starting Analysis of correlation '
    
    concepts_dictionary = msigdb.parser_msigdb(open('/exds/users/oabalbin/downloads/c5.bp.v2.5.symbols.gmt'))
        
    expMat_bfrm_coding = pickle.load(open(outfile1))
    expMat_bfrm_noveltu = pickle.load(open(outfile2))
    #correlation_matrix = pickle.load(open(outfile3))
    correlation_matrix = np.load(outfile4)
    #concepts_dictionary = pickle.load(open(concepts))
        
    if (expMat_bfrm_coding.expVal.shape[0] != correlation_matrix.shape[0]) or (expMat_bfrm_noveltu.expVal.shape[0] != correlation_matrix.shape[1]):
        print expMat_bfrm_coding.expVal.shape
        print expMat_bfrm_noveltu.expVal.shape
        print correlation_matrix.shape        
        sys.exit(0)
    
    target_genes = np.array(expMat_bfrm_coding.geneSym)
    noveltus = np.array(expMat_bfrm_noveltu.geneSym)
    # comparison can be one of three values: less, greater or both
    hugo_names = tga.translate_refflat2hugo(target_genes)
    hugo_names_arr = np.array(hugo_names)
    
    EZH2_targets = set_gene_list(tp.list_of_names(open(seed_file['targets'])))
    EZH2_targets = set_gene_list(tga.translate_hugo2refflat(EZH2_targets))
    
    #outfile7 = open(outfile7,'w')
    #outfile7.write("geneName1"+'\t'+"geneName"+'\t'+",".join(expMat_bfrm_noveltu.geneSym).replace(',','\t')+'\n')
    #write_array2file(outfile7, correlation_matrix, expMat_bfrm_coding.geneSym, EZH2_targets)
    
    print 'TUs Frequency based on threshold '
    threshold= -0.7
    comparison='less'
    EZH2_targets_in_experiment = get_common_targets(set(EZH2_targets),set(target_genes))
    
    # Obtaining list of TUs negatively correlated with genes   
    # Frequent_tusL, frequent_genesL = It is a column array of the TUs and Genes 
    TUsDictL, highly_correlated_ng = get_gene_transcript_association_threshold(correlation_matrix, threshold, noveltus, target_genes, \
                                                                             EZH2_targets_in_experiment, open(outfile5b,'w'),\
                                                                             open(outfile5c,'w'), comparison)
    
    # Obtaining list of TUs positively correlated with genes
    threshold= 0.7
    comparison="greater"
    TUsDictG, highly_correlated_pg = get_gene_transcript_association_threshold(correlation_matrix, threshold, noveltus, target_genes, \
                                                                             EZH2_targets_in_experiment, open(outfile5b,'w'),\
                                                                             open(outfile5c,'w'), comparison)
    
    # Total number of genes negative and positive correlated
    total_negative_genes = len(set(tga.translate_refflat2hugo( target_genes[highly_correlated_ng] )))
    total_positive_genes = len(set(tga.translate_refflat2hugo( target_genes[highly_correlated_pg] )))
    ## Look for TUs that are in both highly anti-correlated with some genes, but highly correlated with other set    
    TUs_of_interest =  set(TUsDictL.keys()).intersection(set(TUsDictG.keys()))   
    TUs_of_interst_dict, TUs_of_interst_dict_len = get_TUs_double_correlated(TUsDictL,TUsDictG,TUs_of_interest)
    
    # Get TUs- genes using the Hugo nomenclature. This is necessary to quantify the enrichment analysis.
    # TUs_of_interest_hugo = [size set 1,size set 2,gene set 1 (-),gene set 2 (+)]
    TUs_of_interest_hugo, TUs_of_interest_len_hugo = get_TUs_double_correlated_hugo(TUs_of_interst_dict,noveltus,target_genes, outfile8)
    
    # Enrichment analysis
    #concept_names=set()
    FDR=0.05
    i=0
    for tui, myvals in TUs_of_interest_hugo.iteritems():
        mykey = noveltus[tui]
        negative_genes, positive_genes = myvals[2], myvals[3]
        num_negative_genes, num_positive_genes = myvals[0], myvals[1]
        
        common_negative_genes_pvalues = np.zeros((len(TUs_of_interest_hugo),len(concepts_dictionary)))
        common_postive_genes_pvalues = np.zeros((len(TUs_of_interest_hugo),len(concepts_dictionary)))
        
        j=0
        for name, concept_genes in concepts_dictionary.iteritems():
            #concept_names.add(name)
            common_negative_genes = len(set(negative_genes).intersection(set(concept_genes)))
            common_postive_genes = len(set(positive_genes).intersection(set(concept_genes)))
            concept_size = len(concept_genes)
            #calc_hypergeometric_test(drawn_white_balls,total_white_balls, total_balls_drawn,total_black_balls, FDR)
            pvalue_ng = calc_hypergeometric_test(common_negative_genes, concept_size, num_negative_genes, (total_negative_genes - num_negative_genes), FDR)
            pvalue_pg = calc_hypergeometric_test(common_postive_genes, concept_size, num_positive_genes, (total_positive_genes - num_positive_genes), FDR)
            common_negative_genes_pvalues[i,j] = pvalue_ng
            common_postive_genes_pvalues[i,j] = pvalue_pg
            #print mykey, name, pvalue_ng, pvalue_pg
            j+=1
        print len(TUs_of_interest_hugo) - i
        
    concept_names=concepts_dictionary.keys()
    concept_TUs_list=[]
    for tt in range(len(TUs_of_interest_hugo)):
        thistu_pvalues = common_negative_genes_pvalues[i,:]
        pqvalues_array = calculate_FDR(thistu_pvalues)
        concept_TUs_list.append(tu_enriched_concepts(noveltus[tt],pqvalues_array, concept_names))
    
    cocept_TU_file = open(dirs[3]+'/'+tstamp+'_concept_TUs_list.txt','w')
    pickle.dump(concept_TUs_list,cocept_TU_file)
    
    
    print len(TUsDictL.keys()), len(TUsDictG.keys()),correlation_matrix.shape[1]
    print len(TUs_of_interest)
    # A matrix of TUs indicators for each gene, and dictionari of TU[ind]=TUs        
    #geneTU_indicator_matrixG, tu_inddictG, All_Corr_TUsG = get_gene_transcript_association(correlation_matrix, toptus, open(outfile5,'w'), 
    #                                                                                    comparison)
    

    sys.exit(0)
    



